from __future__ import print_function

from keras.layers.normalization import BatchNormalization
from keras.models import Sequential, Graph
from keras.layers.convolutional import Convolution2D, MaxPooling2D, ZeroPadding2D
from keras.layers.core import Activation, Dense, Flatten, Dropout, Reshape, RepeatVector
from keras.layers.recurrent import LSTM
from keras.regularizers import l2
from keras import backend as K


def scale(x):
    return (x - K.mean(x)) / K.std(x)
    # return x


def get_model():
    conv = Sequential()
    conv.add(Activation(activation=scale, input_shape=(30, 128, 128)))

    conv.add(Convolution2D(64, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(64, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Convolution2D(128, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(128, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Convolution2D(256, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(256, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(Convolution2D(512, 3, 3, border_mode='same'))
    conv.add(Activation('relu'))
    conv.add(MaxPooling2D(pool_size=(2, 2)))
    conv.add(BatchNormalization())
    conv.add(Dropout(0.2))

    conv.add(Flatten())
    conv.add(Reshape((32, 16)))
    conv.add(LSTM(512, return_sequences=False, input_shape=(32, 16)))
    conv.add(Dropout(0.5))
    #conv.add(Reshape((1, 2048)))

    meta = Sequential()
    meta.add(RepeatVector(32))
    # meta.add(Dense(512, input_dim=4))
    # meta.add(Activation('relu'))
    # meta.add(Reshape((1, 512)))
    meta.add(LSTM(512, input_shape=(32, 4), return_sequences=False))
    meta.add(Dropout(0.5))
    #meta.add(Reshape((1, 512)))

    model = Graph()
    model.add_input(name='conv_input', input_shape=(30, 128, 128))
    model.add_input(name='meta_input', input_shape=(4,))
    model.add_node(conv, name='conv', input='conv_input')
    model.add_node(meta, name='meta', input='meta_input')
    model.add_node(LSTM(512, return_sequences=True), name='merge', inputs=['conv', 'meta'], merge_mode='concat')
    model.add_node(Dropout(0.5), name='lstm_merge_do', input='lstm_merge')
    model.add_node(LSTM(512, return_sequences=False), name='lstm_1', input='lstm_merge_do')
    model.add_node(Dropout(0.5), name='lstm_do_1', input='lstm_1')
    # model.add_node(LSTM(512), name='lstm_1', input='lstm_do_0')
    # model.add_node(Dropout(0.5), name='lstm_do_1', input='lstm_1')
    model.add_node(Dense(1), name='merge_out', input='lstm_do_0')
    model.add_output(name='output', input='merge_out')
    return model
